#include "tracknetv3_trt.h"
#include <iostream>
#include <chrono>
#include <cassert>
#include <vector>
#include <string>
#include <fstream>
#include <iomanip>

class VideoProcessor {
private:
    cv::VideoCapture cap_;
    cv::VideoWriter writer_;
    TrackNetV3_TRT tracknet_;
    std::ofstream csv_file_;
    
    int fps_;
    int width_;
    int height_;
    int input_width_;
    int input_height_;
    float ratio_x_;
    float ratio_y_;
    
    std::vector<cv::Mat> frame_buffer_;
    int frame_count_;
    int processed_count_;
    
public:
    VideoProcessor() : frame_count_(0), processed_count_(0) {}
    
    bool initialize(const std::string& video_path, const std::string& engine_path, 
                   const std::string& output_video_path, const std::string& output_csv_path) {
        // 打开输入视频
        cap_.open(video_path);
        if (!cap_.isOpened()) {
            std::cerr << "Error: Cannot open video file: " << video_path << std::endl;
            return false;
        }
        
        // 获取视频属性
        fps_ = static_cast<int>(cap_.get(cv::CAP_PROP_FPS));
        width_ = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_WIDTH));
        height_ = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_HEIGHT));
        int total_frames = static_cast<int>(cap_.get(cv::CAP_PROP_FRAME_COUNT));
        
        std::cout << "Video Info: " << width_ << "x" << height_ << ", FPS: " << fps_ 
                  << ", Total frames: " << total_frames << std::endl;
        
        // 初始化TensorRT模型
        std::cout << "Initializing TrackNetV3 TensorRT engine..." << std::endl;
        if (!tracknet_.init(engine_path)) {
            std::cerr << "Failed to initialize TrackNetV3" << std::endl;
            return false;
        }
        
        // 获取模型输入尺寸并计算缩放比例
        int input_channels;
        tracknet_.getInputSize(input_width_, input_height_, input_channels);
        ratio_x_ = static_cast<float>(width_) / input_width_;
        ratio_y_ = static_cast<float>(height_) / input_height_;
        
        std::cout << "Model input size: " << input_width_ << "x" << input_height_ 
                  << ", channels: " << input_channels << std::endl;
        std::cout << "Scaling ratio: " << ratio_x_ << "x" << ratio_y_ << std::endl;
        
        // 创建输出视频
        int fourcc = cv::VideoWriter::fourcc('m', 'p', '4', 'v');
        writer_.open(output_video_path, fourcc, fps_, cv::Size(width_, height_));
        if (!writer_.isOpened()) {
            std::cerr << "Error: Cannot create output video: " << output_video_path << std::endl;
            return false;
        }
        
        // 创建CSV文件
        csv_file_.open(output_csv_path);
        if (!csv_file_.is_open()) {
            std::cerr << "Error: Cannot create CSV file: " << output_csv_path << std::endl;
            return false;
        }
        
        // 写入CSV表头
        csv_file_ << "Frame,Visibility,X,Y\n";
        
        std::cout << "Video processor initialized successfully" << std::endl;
        return true;
    }
    
    void processVideo() {
        std::cout << "Start processing video..." << std::endl;
        
        auto total_start = std::chrono::high_resolution_clock::now();
        int batch_size = 1; // 每次处理一个三帧序列
        
        cv::Mat frame;
        while (true) {
            // 读取帧直到凑齐三帧
            while (frame_buffer_.size() < 3) {
                bool success = cap_.read(frame);
                if (!success) {
                    break;
                }
                frame_buffer_.push_back(frame.clone());
                frame_count_++;
            }
            
            // 如果没有足够的帧继续处理，退出循环
            if (frame_buffer_.size() < 3) {
                std::cout << "Not enough frames to continue processing. Buffer size: " 
                          << frame_buffer_.size() << std::endl;
                break;
            }
            
            // 处理当前三帧序列
            if (!processFrameSequence()) {
                std::cerr << "Error processing frame sequence starting at frame " 
                          << (frame_count_ - frame_buffer_.size() + 1) << std::endl;
                break;
            }
            
            // 保留最后一帧作为下一序列的第一帧
            if (frame_buffer_.size() >= 3) {
                // 保留最后两帧，与下一帧组成新的三帧序列
                std::vector<cv::Mat> new_buffer;
                new_buffer.push_back(frame_buffer_[1]);
                new_buffer.push_back(frame_buffer_[2]);
                frame_buffer_ = new_buffer;
            }
            
            // 显示进度
            if (processed_count_ % 100 == 0) {
                std::cout << "Processed " << processed_count_ << " frames..." << std::endl;
            }
        }
        
        auto total_end = std::chrono::high_resolution_clock::now();
        auto total_duration = std::chrono::duration_cast<std::chrono::seconds>(total_end - total_start);
        
        std::cout << "Video processing completed!" << std::endl;
        std::cout << "Total frames processed: " << processed_count_ << std::endl;
        std::cout << "Total time: " << total_duration.count() << " seconds" << std::endl;
    }
    
private:
    bool processFrameSequence() {
        if (frame_buffer_.size() < 3) {
            return false;
        }
        
        // 提取当前三帧序列
        std::vector<cv::Mat> current_frames(frame_buffer_.begin(), frame_buffer_.begin() + 3);
        
        // 执行推理
        std::vector<cv::Mat> output_heatmaps;
        auto start = std::chrono::high_resolution_clock::now();
        
        if (!tracknet_.process(current_frames, output_heatmaps)) {
            std::cerr << "Inference failed for frame sequence starting at frame " 
                      << (frame_count_ - frame_buffer_.size() + 1) << std::endl;
            return false;
        }
        
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        
        // 处理每一帧的输出
        for (int i = 0; i < output_heatmaps.size() && i < 3; ++i) {
            int current_frame_index = frame_count_ - frame_buffer_.size() + i + 1;
            processSingleFrame(current_frames[i], output_heatmaps[i], current_frame_index);
            processed_count_++;
        }
        
        if (processed_count_ % 10 == 0) {
            std::cout << "Inference time for sequence: " << duration.count() << " ms" << std::endl;
        }
        
        return true;
    }
    
    void processSingleFrame(const cv::Mat& original_frame, const cv::Mat& heatmap, int frame_number) {
        cv::Mat output_frame = original_frame.clone();
        
        // 从热力图中提取羽毛球位置
        int ball_x, ball_y;
        bool ball_detected = false;
        
        // 处理多通道热力图（取第一个通道或合并所有通道）
        cv::Mat processed_heatmap;
        if (heatmap.channels() > 1) {
            // 将多通道热力图合并或取第一个通道
            std::vector<cv::Mat> channels;
            cv::split(heatmap, channels);
            processed_heatmap = channels[0]; // 取第一个通道
        } else {
            processed_heatmap = heatmap;
        }
        
        ball_detected = tracknet_.getBallCenter(processed_heatmap, ball_x, ball_y, 0.5f);
        
        // 将坐标缩放回原始图像尺寸
        int scaled_x = static_cast<int>(ball_x * ratio_x_);
        int scaled_y = static_cast<int>(ball_y * ratio_y_);
        
        // 写入CSV
        int visibility = ball_detected ? 1 : 0;
        csv_file_ << frame_number << "," << visibility << "," << scaled_x << "," << scaled_y << "\n";
        
        // 在图像上绘制检测结果
        if (ball_detected) {
            // 绘制羽毛球位置
            cv::circle(output_frame, cv::Point(scaled_x, scaled_y), 10, cv::Scalar(0, 255, 0), -1);
            cv::circle(output_frame, cv::Point(scaled_x, scaled_y), 10, cv::Scalar(255, 255, 255), 2);
            
            // 添加帧编号和检测信息
            std::string info_text = "Frame: " + std::to_string(frame_number) + 
                                   " Ball: (" + std::to_string(scaled_x) + ", " + std::to_string(scaled_y) + ")";
            cv::putText(output_frame, info_text, cv::Point(10, 30), 
                       cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(255, 255, 255), 2);
        } else {
            // 没有检测到球
            std::string info_text = "Frame: " + std::to_string(frame_number) + " No ball detected";
            cv::putText(output_frame, info_text, cv::Point(10, 30), 
                       cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 0, 255), 2);
        }
        
        // 写入输出视频
        writer_.write(output_frame);
    }
    
public:
    void cleanup() {
        if (cap_.isOpened()) {
            cap_.release();
        }
        if (writer_.isOpened()) {
            writer_.release();
        }
        if (csv_file_.is_open()) {
            csv_file_.close();
        }
        std::cout << "Resources cleaned up" << std::endl;
    }
};

int main(int argc, char** argv) {
    if (argc < 3) {
        std::cout << "Usage: " << argv[0] << " <input_video> <engine_file> [output_dir]" << std::endl;
        std::cout << "Example: " << argv[0] << " test_video.mp4 tracknetv3.engine ./output" << std::endl;
        return -1;
    }
    
    std::string video_path = argv[1];
    std::string engine_path = argv[2];
    std::string output_dir = (argc > 3) ? argv[3] : "./output";
    
    // 创建输出目录
    system(("mkdir -p " + output_dir).c_str());
    
    // 生成输出文件路径
    size_t last_slash = video_path.find_last_of("/\\");
    size_t last_dot = video_path.find_last_of(".");
    std::string video_name = (last_slash != std::string::npos) ? 
                            video_path.substr(last_slash + 1, last_dot - last_slash - 1) : 
                            video_path.substr(0, last_dot);
    
    std::string output_video_path = output_dir + "/" + video_name + "_pred.mp4";
    std::string output_csv_path = output_dir + "/" + video_name + "_ball.csv";
    
    std::cout << "=== TrackNetV3 Video Processing ===" << std::endl;
    std::cout << "Input video: " << video_path << std::endl;
    std::cout << "Engine file: " << engine_path << std::endl;
    std::cout << "Output video: " << output_video_path << std::endl;
    std::cout << "Output CSV: " << output_csv_path << std::endl;
    
    VideoProcessor processor;
    
    try {
        if (!processor.initialize(video_path, engine_path, output_video_path, output_csv_path)) {
            std::cerr << "Failed to initialize video processor" << std::endl;
            return -1;
        }
        
        processor.processVideo();
        processor.cleanup();
        
        std::cout << "=== Processing Completed Successfully ===" << std::endl;
        std::cout << "Results saved to: " << output_dir << std::endl;
        
    } catch (const std::exception& e) {
        std::cerr << "Exception caught: " << e.what() << std::endl;
        processor.cleanup();
        return -1;
    } catch (...) {
        std::cerr << "Unknown exception caught" << std::endl;
        processor.cleanup();
        return -1;
    }
    
    return 0;
}
